#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <cv.h>
#include <highgui.h>
#include <ctime>
#include <vector>
#include "opencv_helper.h"
#include "cvBlobsLib/BlobResult.h"
#include "cvBlobsLib/blob.h"

using namespace std;

// TODO: convert to using C++ cv:: libraries, instead of C ones

// Global processing variables
IplImage *hsv;
IplImage *smooth;
IplImage *smoothH; 
IplImage *smoothS; 
IplImage *smoothV; 
IplImage *sobel;
IplImage *sobelGray;
IplImage *canny;
IplImage *edges;
IplImage *pyrseg;
CvMemStorage *storage;
CvSeq *components;
IplImage *imgOutBinary;
IplImage *imgOutErode;

typedef struct {
  int hwraps;
  int minH;
  int maxH;
  int minS;
  int maxS;
  int minV;
  int maxV;
} thresh_t;

inline int min(int x, int y) {
    return (x < y) ? x : y;
}

inline int max(int x, int y) {
    return (x >= y) ? x : y;
}

// Small chaining functions
IplImage *addConeModels(char *filename);

// Core processing functions
int initGlobals(int width, int height, int channels);
int initImageProcessing(char **coneModelFiles, thresh_t *thresh);
int processImage(IplImage *imgOut, IplImage *imgIn, int height, int width, int channels, thresh_t *thresh);



int initGlobals(int width, int height, int channels)
{
  hsv = cvCreateImage(cvSize(width, height), 8, channels); 
  smooth = cvCreateImage(cvSize(width, height), 8, channels); 
  smoothH = cvCreateImage(cvSize(width, height), 8, 1); 
  smoothS = cvCreateImage(cvSize(width, height), 8, 1); 
  smoothV = cvCreateImage(cvSize(width, height), 8, 1); 
  sobel = cvCreateImage(cvSize(width, height), IPL_DEPTH_16S, 1); 
  sobelGray = cvCreateImage(cvSize(width, height), 8, 1); 
  canny = cvCreateImage(cvSize(width, height), 8, 1); 
  edges = cvCreateImage(cvSize(width, height), 8, 1); 
  pyrseg = cvCreateImage(cvSize(width, height), 8, 1); 
  storage = cvCreateMemStorage(0);
  components = NULL;

  imgOutBinary = cvCreateImage(cvSize(width, height), IPL_DEPTH_8U, 1);
  imgOutErode = cvCreateImage(cvSize(width, height), IPL_DEPTH_8U, 1);

}

IplImage *addConeModels(char **filenames)
{
  IplImage *unified = cvCreateImage(cvSize(1000, 1), IPL_DEPTH_8U, 3);
  int unified_pxl = 0;

  while (*filenames) {
    IplImage *img = cvLoadImage(*filenames);
    ImgWrap8u3c imgw(img);
 
    char c0;
    char c1;
    char c2;
    for (int x = 0; x < imgw->width; x++) {
      for (int y = 0; y < imgw->height; y++) {
        c0 = imgw[y][x].c0;
        c1 = imgw[y][x].c1;
        c2 = imgw[y][x].c2;

        // Copy non-zero pixels
        if (c0 != 0 || c1 != 0 || c2 != 0) {
          // Grow if necessary
          if (unified_pxl == unified->width) {
            IplImage *new_uni = cvCreateImage(cvSize(unified_pxl*2, 1), IPL_DEPTH_8U, 3);
            ImgWrap8u3c new_uniw(new_uni);
            ImgWrap8u3c unifiedw(unified);

            for (int i = 0; i < unified_pxl; i++) {
              new_uniw[0][i].c0 = unifiedw[0][i].c0;
              new_uniw[0][i].c1 = unifiedw[0][i].c1;
              new_uniw[0][i].c2 = unifiedw[0][i].c2;
            }
            cvReleaseImage(&unified);
            unified = new_uni;
          }

          ImgWrap8u3c unifiedw(unified);
    
          // Copy pixels
          unifiedw[0][unified_pxl].c0 = c0;
          unifiedw[0][unified_pxl].c1 = c1;
          unifiedw[0][unified_pxl].c2 = c2;
          unified_pxl++;
        }
      }
    }

    cvReleaseImage(&img);

    ++filenames;
  }

  return unified;
}

int initImageProcessing(char **coneModelFiles, thresh_t *thresh)
{
  IplImage *unifiedImg = addConeModels(coneModelFiles);
  IplImage *unifiedHsv = cvCreateImage(cvSize(unifiedImg->width, 1), IPL_DEPTH_8U, 3);
  cvCvtColor(unifiedImg, unifiedHsv, CV_BGR2HSV);
  ImgWrap8u3c uniw(unifiedHsv);

  // Calculate the mean
  int accumH = 0;
  int accumH2 = 0;
  int accumS = 0;
  int accumV = 0;
  int hshift;
  int samples = uniw->width;

  for (int x = 0; x < samples; x++) {
      // TODO: double check this
      hshift = uniw[0][x].c0;
      if (hshift < 128) {
        hshift += 256;
      }
      accumH += uniw[0][x].c0;
      accumH2 += hshift;
      accumS += uniw[0][x].c1;
      accumV += uniw[0][x].c2;
  }

  int meanH = accumH / samples;
  int meanH2 = accumH2 / samples; // 128 - 384
  int meanS = accumS / samples;
  int meanV = accumV / samples;

  // Calculate the stddev 
  accumH = 0;
  accumH2 = 0;
  accumS = 0;
  accumV = 0;
  for (int x = 0; x < samples; x++) {
    hshift = uniw[0][x].c0;
    if (hshift < 128) {
      hshift += 256;
    }
    accumH += (uniw[0][x].c0 - meanH) * (uniw[0][x].c0 - meanH);
    accumH2 += (hshift - meanH2) * (hshift - meanH2);
    accumS += (uniw[0][x].c1 - meanS) * (uniw[0][x].c1 - meanS);
    accumV += (uniw[0][x].c2 - meanV) * (uniw[0][x].c2 - meanV);
  }

  int varH = accumH / samples;
  int varH2 = accumH2 / samples;
  int varS = accumS / samples;
  int varV = accumV / samples;
  
  int stdvH = sqrt(varH);
  int stdvH2 = sqrt(varH2);
  int stdvS = sqrt(varS);
  int stdvV = sqrt(varV);

  int coef = 2;

  if (stdvH <= stdvH2) {
    if (2*coef*stdvH > 255) {
      thresh->minH = 0;
      thresh->maxH = 255;
    }
    else {
      thresh->minH = meanH - coef*stdvH;
      thresh->maxH = meanH + coef*stdvH;
    }
  }
  else {
    // Need to shift appropriately
    if (meanH2 > 255) {
      meanH2 -= 256;
    }
    printf("mean: %d\n", meanH2);

    if (2*coef*stdvH2 > 255) {
      thresh->minH = 0;
      thresh->maxH = 255;
    }
    else {
      thresh->minH = meanH2 - coef*stdvH2;
      thresh->maxH = meanH2 + coef*stdvH2;
    }
    printf("thresh: %d %d\n", thresh->minH, thresh->maxH);
  }

  if (thresh->minH < 0) {
    thresh->minH += 256;
    thresh->maxH += 256;
    thresh->hwraps = 1;
  }
  else {
    thresh->hwraps = 0;
  }

  thresh->minS = meanS - coef*stdvS;
  thresh->maxS = meanS + coef*stdvS;
  thresh->minV = meanV - coef*stdvV;
  thresh->maxV = meanV + coef*stdvV;

  return 0;

  // TODO
  // Read in image set, remove all white pixels and put all other pixels into a common array
  // Determine 3-d HSV distribution of points (calculated mean + stdev in each axis)
  // Set threshold values based on the data (make the thresholds generous so we overgrab -- we'll prune later)
}

int destroyGlobals()
{
  cvReleaseImage(&hsv);
  cvReleaseImage(&smooth);
  cvReleaseImage(&smoothH);
  cvReleaseImage(&smoothS);
  cvReleaseImage(&smoothV);
  cvReleaseImage(&sobel);
  cvReleaseImage(&sobelGray);
  cvReleaseImage(&canny);
  cvReleaseImage(&edges);
  cvReleaseImage(&pyrseg);
  cvReleaseMemStorage(&storage);
  if (components != NULL) {
      printf("how do I release this memory?\n");
  }
  cvReleaseImage(&imgOutBinary);
  cvReleaseImage(&imgOutErode);
}



int processImage(IplImage *imgOut, IplImage *imgIn, int height, int width, int channels, thresh_t *thresh)
{
  ////////////////////////////////////////////////////
  //double scale, shift;
  //double minVal, maxVal;

  //cvCvtColor(imgIn, hsv, CV_BGR2HSV);
  //cvSmooth(hsv, smooth, CV_MEDIAN, 3, 0, 0);
  //cvSplit(smooth, smoothS, smoothH, smoothV, NULL);
  //
  //cvSobel(smoothH, sobel, 1, 0, 5);

  //cvMinMaxLoc(sobel, &minVal, &maxVal);
  //scale = 255/(0.75 * (maxVal - minVal));
  //shift = -minVal*scale;
  //cvConvertScale(sobel, sobelGray, scale, shift);

  //cvCanny(smoothH, edges, 110, 170, 3);
  //cvPyrSegmentation(smoothH, pyrseg, storage, &components, 3, 120, 140);
  //cvMerge(pyrseg, edges, edges, NULL, imgOut);
  ////////////////////////////////////////////////////

  // TODO
  // Do some basic filtering to clean the image up a bit
  // Remove everything not within the 3-d HSV thresholds determined earlier
  // Filter out regions that have low density of pixels left
  // Do some basic shape analysis of the blobs
  // Assign a confidence to each "cone"
  // Later: Add some hysteresis between frames, to prevent cones popping around -- confidence will grow or shrink with time
  // Add some feedback mechanism to train the algorithm and thresholds

  // TODO: color by how close it is to the mean
  cvCvtColor(imgIn, hsv, CV_BGR2HSV);
  ImgWrap8u3c hsvw(hsv); 
  ImgWrap8u3c outw(imgOut); 
  ImgWrap8u1c binw(imgOutBinary);
  for (int x = 0; x < hsvw->width; x++) {
    for (int y = 0; y < hsvw->height; y++) {
      int h = hsvw[y][x].c0;
      int s = hsvw[y][x].c1;
      int v = hsvw[y][x].c2;
      int htest;

      if (thresh->hwraps) {
        h += 256;
      }

      if (h <= thresh->maxH && h >= thresh->minH &&
          s <= thresh->maxS && s >= thresh->minS &&
          v <= thresh->maxV && v >= thresh->minV) {
        binw[y][x].c0 = 255; 
        //outw[y][x].c0 = h;
        //outw[y][x].c1 = s;
        //outw[y][x].c2 = v;
      }
      else {
        binw[y][x].c0 = 0; 
        //outw[y][x].c0 = 0;
        //outw[y][x].c1 = 0;
        //outw[y][x].c2 = 0;
      }
    }
  }

  // TODO
  // Not sure if this speeds up the blob filtering below or not, maybe just
  // rely on that.  Also, may want to do a dilate then erode to join close
  // objects
  cvErode(imgOutBinary, imgOutErode, NULL, 2);
  cvDilate(imgOutErode, imgOutBinary, NULL, 2);

  // Blob detection
  CBlobResult blobs;
  blobs = CBlobResult(imgOutBinary, NULL, 0);
  int blob_min_size = 300;
  blobs.Filter(blobs, B_INCLUDE, CBlobGetArea(), B_GREATER, blob_min_size);
  int num_blobs = blobs.GetNumBlobs();

  CBlob *blob;
  CvRect bb;
  CvPoint p1;
  CvPoint p2;

  // Shape filtering
  float top_chop = .1;
  float bot_chop = .3;
  
  ImgWrap8u1c blobimg_wrap(imgOutBinary);
  int top_left_x;
  int top_left_y;
  int bot_left_y;
  for (int b = 0; b < num_blobs; b++) {
    blob = blobs.GetBlob(b);
    bb = blob->GetBoundingBox();

    //p1 = cvPoint(bb.x, bb.y);
    //p2 = cvPoint(bb.x + bb.width, bb.y + bb.height);
    //cvRectangle(imgOutBinary, p1, p2, CV_RGB(192, 0, 0), 2, 8, 0);
    //cvRectangle(imgIn, p1, p2, CV_RGB(192, 0, 0), 2, 8, 0);

    top_left_x = bb.x;
    top_left_y = bb.y + (bb.height * top_chop);
    bot_left_y = bb.y + (bb.height * (1 - bot_chop));

    // TODO
    // Get moments
    // Pre-filter objects without centered verticle moment, or cone-like horizontal moment
    // Also filter by second moments?

    int x,y;
    int point_count = bot_left_y - top_left_y;
    vector<cv::Point> *points = new vector<cv::Point>(point_count);
    for (y = top_left_y; y < bot_left_y; y++) {
      x = bb.x;
      for (x = bb.x; blobimg_wrap[y][x].c0 == 0; x++)
        // Hunt for the first edge from the edge of the bounding box
        ;

      points->push_back(cv::Point(x, y));
    }

    // TODO: also hunt on right side for other line

    cv::Mat point_mat = cv::Mat(*points);
    cv::Vec4f best_line;
    fitLine(point_mat, best_line, CV_DIST_L2, 0, 0.1, 0.1);

    // TODO: draw

    // TODO: Make sure lines on both side are symmetric, slanted the right direction at a good
    // angle, and are balanced on the center moment

    delete points;
  }

  return 0;
}



// Main loop
int main(int argc, char *argv[])
{
  CvCapture *capture;
  IplImage *img;
  IplImage *procImg;
  int height, width, channels;
  int frames;
  time_t last, now;

  // Create the windows
  cvNamedWindow("procWin", CV_WINDOW_AUTOSIZE); 
  cvNamedWindow("origWin", CV_WINDOW_AUTOSIZE); 
  cvMoveWindow("procWin", 10, 10);
  cvMoveWindow("origWin", 700, 400);

  // Choose the source
  //if (argc >= 2) {
    capture = cvCaptureFromFile(argv[1]);
  //} else {
  //  capture = cvCaptureFromCAM(0);
  //}

  // Grab the first frame
  cvGrabFrame(capture);
  img = cvRetrieveFrame(capture);
  if (!img) {
    printf("Could not load image\n");
    exit(0);
  }
  
  // Get the image parameters
  width     = img->width;
  height    = img->height;
  channels  = img->nChannels;
  printf("Height: %d, Width: %d, Channels: %d\n", height, width, channels);

  // Video writer
  // TODO
  //CvVideoWriter *video_writer = cvCreateVideoWriter("/Users/tnorrie/Desktop/results.avi", CV_FOURCC('P', 'I', 'M', '1'), 25, cvSize(width, height), 1);

  // Init globals
  initGlobals(width, height, channels);
  procImg = cvCreateImage(cvSize(width, height), 8, channels); 

  // Initialization stuff
  // TODO
  thresh_t thresh;
  initImageProcessing(&argv[2], &thresh);

  // Keep track of elapsed time
  now = time(0);
  last = now;

  while (cvGrabFrame(capture)) {
    // Grab the image
    img = cvRetrieveFrame(capture);
    if (!img) {
      printf("Could not load image\n");
      exit(0);
    }

    // Process the image
    // TODO
    processImage(procImg, img, height, width, channels, &thresh);
    
    // Show the processing results
    cvShowImage("origWin", img);
    cvShowImage("procWin", imgOutBinary);//procImg);// TODO

    // Write results to a file
    // TODO
    //cvWriteFrame(video_writer, procImg);

    // Keep track of FPS
    now = time(0);
    frames++;
    if (now - last >= 5) {
        printf("%.1f FPS\n", frames/5.0);
        frames = 0;
        last = now;
    }

    // Yield for GUI events
    cvWaitKey(1);
  }

  // Clean up
  //cvReleaseVideoWriter(&video_writer); // TODO
  destroyGlobals();
  cvReleaseImage(&img);
  cvReleaseImage(&procImg);
  cvReleaseCapture(&capture);
  cvDestroyWindow("mainWin");
  cvDestroyWindow("origWin");

  return 0;
}


